using System;
using System.Collections.Generic;
using System.Text;
using DotNetMatrix;

namespace Blimpbots.Recognition.Algorithms
{
    public class PointMassKalmanFilter : KalmanFilter
    {
        private double _last_t;
        double _stDev_x, _stDev_y, _stDev_z;

        public PointMassKalmanFilter(double stDev_x, double stDev_y,
                                double stDev_z, double init_stDev, Point3D x, double t)
        {
            _F = GeneralMatrix.Identity(6, 6);

            _stDev_x = stDev_x;
            _stDev_y = stDev_y;
            _stDev_z = stDev_z;

            _last_t = t;

            // Note that we have no certainty about the position
            _P = GeneralMatrix.Identity(6, 6).MultiplyEquals(init_stDev * init_stDev);

            _x = new GeneralMatrix(new double[] { x.X, x.Y, x.Z, 0, 0, 0 }, 6);

        }

        public Point3D Predict(double t)
        {
            double dT = t - _last_t;

            #region Prepare Q
            _Q = CalcQ(dT);
            #endregion

            #region Prepare F

            // Position estimate is the same as before
            //   plus velocity times dT
            _F.Array[0][3] = dT;
            _F.Array[1][4] = dT;
            _F.Array[2][5] = dT;

            #endregion

            #region Predict and Prepare H
            GeneralMatrix x = PredictNonDestructive();
            #endregion

            #region Convert X to Point3D
            Point3D result = new Point3D(x.Array[0][0], x.Array[1][0], x.Array[2][0]);
            #endregion

            return result;
        }

        public Point3D Update(Point3D z, double meas_stDev_x, double meas_stDev_y,
                              double meas_stDev_z, GeneralMatrix ellipseRot, double t)
        {
            #region Convert Z to Matrix
            GeneralMatrix zMat = new GeneralMatrix(new double[] { z.X, z.Y, z.Z }, 3);
            #endregion

            double dT = t - _last_t;
            _last_t = t;

            #region Prepare Q
            _Q = CalcQ(dT);
            #endregion

            #region Prepare F

            // Position estimate is the same as before
            //   plus velocity times dT
            _F.Array[0][3] = dT;
            _F.Array[1][4] = dT;
            _F.Array[2][5] = dT;

            #endregion

            #region Predict and Prepare H

            GeneralMatrix B = new GeneralMatrix(3, 6);
            B.Array[0][0] = 1;
            B.Array[1][1] = 1;
            B.Array[2][2] = 1;
            _H = B;


            GeneralMatrix _xHat = Predict();

            #endregion

            #region Prepare R

            _R = GeneralMatrix.Identity(3, 3);
            _R.Array[0][0] = meas_stDev_x * meas_stDev_x;
            _R.Array[1][1] = meas_stDev_y * meas_stDev_y;
            _R.Array[2][2] = meas_stDev_z * meas_stDev_z;

            // S_2 is the diagonal error matrix inverted, and that result squared
            //  i.e. inv(diag([sigma_x, sigma_y, sigma_z]))^2

            GeneralMatrix S_2 = GeneralMatrix.Identity(3,3);
            S_2.Array[0][0] = 1.0 / (meas_stDev_x * meas_stDev_x);
            S_2.Array[1][1] = 1.0 / (meas_stDev_y * meas_stDev_y);
            S_2.Array[2][2] = 1.0 / (meas_stDev_z * meas_stDev_z);

            // From SVD paper: R = inv(R*S^(-2)*Rt)
            //_R = ellipseRot * S_2 * ellipseRot.Transpose();
            //_R = _R.Inverse();

            #endregion

            #region Update
            GeneralMatrix x = Update(zMat);
            #endregion

            // TODO: normalize?
            #region Convert Z to Point3D
            Point3D result = new Point3D(x.Array[0][0], x.Array[1][0], x.Array[2][0]);
            #endregion

            return result;
        }

        private GeneralMatrix CalcQ(double dT)
        {
            //GeneralMatrix G = new GeneralMatrix(new double[6][] {
            //   new double[1] { (dT*dT)/2 }, // x
            //   new double[1] { (dT*dT)/2 }, // y
            //   new double[1] { (dT*dT)/2 }, // z
            //   new double[1] { dT },        // x dot
            //   new double[1] { dT },        // y dot
            //   new double[1] { dT }         // z dot
            //   });

            //GeneralMatrix Sigma = new GeneralMatrix(new double[] {
            //   _stDev_x,
            //   _stDev_y,
            //   _stDev_z,
            //   _stDev_x,
            //   _stDev_y,
            //   _stDev_z,
            //   }, 1);

            //Q = G * Sigma;
            //Q = Q * Q.Transpose();

            GeneralMatrix Q = GeneralMatrix.Identity(6, 6);

            Q.Array[0][0] = _stDev_x * (dT * dT) / 2;
            Q.Array[1][1] = _stDev_y * (dT * dT) / 2;
            Q.Array[2][2] = _stDev_z * (dT * dT) / 2;
            Q.Array[3][3] = _stDev_x * dT;
            Q.Array[4][4] = _stDev_y * dT;
            Q.Array[5][5] = _stDev_z * dT;

            return Q;
        }
    }
}
